package com.tool.grab.vpn.tunnel;

import android.net.VpnService;

import com.tool.grab.Constant;
import com.tool.grab.vpn.header.IPPacket;
import com.tool.grab.vpn.header.UDPPacket;
import com.tool.grab.vpn.task.TunnelEventListener;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.DatagramChannel;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;

public class UDPTunnel implements TunnelEventListener {
    private DatagramChannel channel;
    private IPPacket packet;

    private int localPort;
    private String localIPString;

    private int remotePort;
    private String remoteIPString;

    private int sendBytes;
    private int receiveBytes;

    public UDPTunnel(Selector selector,IPPacket packet) throws IOException{
        this.packet = packet;
        UDPPacket udpPacket = new UDPPacket(packet.getBytes(),packet.getOffset() + packet.getHeaderLength(),packet.getValidLength());
        localPort = udpPacket.getLocalPort();
        localIPString = packet.getLocalIpString();
        remotePort = udpPacket.getRemotePort();
        remoteIPString = packet.getRemoteIpString();

        channel = DatagramChannel.open();
        channel.configureBlocking(false);
        channel.socket().bind(new InetSocketAddress(localIPString,localPort));
        channel.register(selector,SelectionKey.OP_CONNECT | SelectionKey.OP_READ,this);
    }



    /**
     * 建立套接字通道 并加入VPN保护
     * @param service
     * @throws IOException
     */
    public void connect(VpnService service) throws IOException{
        if(isOpen()){
            if(service.protect(channel.socket())){
                channel.connect(new InetSocketAddress(remoteIPString,remotePort));
            }
        }
    }

    @Override
    public void keyReady(SelectionKey key) throws IOException {
        if(key.isConnectable()){
            onWrite(packet);
        } else if (key.isReadable()) {
            onRead();
        }
    }

    public void onWrite(IPPacket packet) throws IOException{
        ByteBuffer buffer = ByteBuffer.wrap(packet.getBytes(),packet.getOffset(),packet.getValidLength());

        while (buffer.hasRemaining()){
            sendBytes += channel.write(buffer);
        }
    }

    private void onRead() throws IOException{
        int readBytes = 0;
        int sum = 0;
        ByteBuffer buffer = ByteBuffer.allocate(Constant.MUTE_SIZE);

        while (readBytes != -1){
            readBytes = channel.read(buffer);
            if(readBytes > 0){
                sum += readBytes;
            }
        }

        if(readBytes <= 0) return;

        receiveBytes += sum;

        byte[] bytes = buffer.array();
    }


    public void onClose() {
        if(channel != null){
            if(channel.isConnected()){
                try {
                    channel.disconnect();
                } catch (IOException e) {}finally {}
            }

            try {
                channel.close();
            } catch (IOException e) {}finally {
                channel = null;
            }
        }
    }

    public int getSendBytes() {
        return sendBytes;
    }

    public int getReceiveBytes() {
        return receiveBytes;
    }

    public boolean isRun() {
        return channel != null && channel.isConnected();
    }

    public boolean isOpen(){
        return channel != null && channel.isOpen();
    }
}
